import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

"""

缺失值处理，处理的使programs的缺失值，但是没有处理全，将缺失值赋值为0

"""

# 读取数据
file_path = '2025_Problem_C_Data/summerOly_programs.csv'
programs = pd.read_csv(file_path, encoding='latin1', sep=':', engine='python')

# 2. 将空字符串替换为 NaN
programs.replace("", np.nan, inplace=True)

# 3. 使用0填充所有缺失值
programs.fillna(0, inplace=True)

# 将数据转换为长格式
programs_long = programs.melt(id_vars=['Sport', 'Discipline', 'Code', 'Sports Governing Body'],
                              var_name='Year',
                              value_name='Medals')

# 处理年份，使其成为整数类型
programs_long['Year'] = programs_long['Year'].apply(lambda x: int(x) if x != 'Total events' else x)

# 绘制处理前后的奖牌分布热力图
plt.figure(figsize=(12, 8))
sns.heatmap(programs.pivot('Sport', 'Year', 'Medals'), annot=True, fmt="d", cmap='viridis')
plt.title('Medals Distribution Before Processing')
plt.savefig('before_processing.png')
plt.show()

plt.figure(figsize=(12, 8))
sns.heatmap(programs_long.pivot('Sport', 'Year', 'Medals'), annot=True, fmt="d", cmap='viridis')
plt.title('Medals Distribution After Processing')
plt.savefig('after_processing.png')
plt.show()

# 绘制特定项目随时间的变化趋势
sport = 'Aquatics'
discipline = 'Swimming'
project_data = programs_long[(programs_long['Sport'] == sport) & (programs_long['Discipline'] == discipline)]

plt.figure(figsize=(10, 6))
sns.lineplot(x='Year', y='Medals', data=project_data, marker='o')
plt.title(f'Medals Trend for {sport} - {discipline}')
plt.xlabel('Year')
plt.ylabel('Medals')
plt.xticks(rotation=45)
plt.grid(True)
plt.savefig(f'trend_{sport}_{discipline}.png')
plt.show()

# 绘制处理前后数据集的差异条形图
years = programs.columns[4:]  # 假设从第5列开始是年份列
sport = 'Aquatics'
discipline = 'Swimming'

plt.figure(figsize=(12, 6))
programs[(programs['Sport'] == sport) & (programs['Discipline'] == discipline)][years].sum().plot(kind='bar', color='blue', label='Before Processing')
programs_long[(programs_long['Sport'] == sport) & (programs_long['Discipline'] == discipline)][years].sum().plot(kind='bar', color='orange', label='After Processing')
plt.title(f'Total Medals for {sport} - {discipline} Before and After Processing')
plt.xlabel('Year')
plt.ylabel('Total Medals')
plt.xticks(rotation=45)
plt.legend()
plt.grid(True)
plt.savefig(f'total_medals_{sport}_{discipline}.png')
plt.show()